#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
SISO UAV Rate PID Control Example using CasADi
@Author: Gordon Liang
"""

import numpy as np
import matplotlib.pyplot as plt
import casadi as ca

from functions.iris_EOM import IrisEOM
from functions.signal_generator import SignalGenerator
from functions.pid_controller import PIDController
from functions.create_integrator import CreateIntegrator

def main():
    ts = 0.0025  # Sample timr
    duration = 1.0  # simulation duration
    iteration = int(duration / ts)
    
    system = IrisEOM()
    signal_gen = SignalGenerator()
    pid_controller = PIDController(kp=50.0, ki=0.5, kd=0.3, u_max=1.0, u_min=-1.0)
    
    integrator = CreateIntegrator(system, ts)
    
    # 初始化状态
    xk = np.array([0.0, 0.0])
    
    # 初始化存储变量
    time = np.zeros(iteration)
    yd = np.zeros(iteration)   # 参考信号
    y = np.zeros(iteration)    # 系统输出
    e = np.zeros(iteration)    # 误差
    u = np.zeros(iteration)    # 控制输入
    para = np.zeros(iteration) # 参数
    
    noise = np.zeros(iteration)
    
    
    # 仿真循环
    for k in range(iteration):
        # 记录时间
        time[k] = k * ts
        
        # 生成参考信号
        yd[k] = 0.5 * signal_gen.generate_reference(time[k])
        
        # 系统仿真一步（使用CasADi积分器）
        xk_ca = ca.DM(xk)
        para_prev = para[k-1] if k > 0 else 0.0
        xk_next = integrator(xk_ca, para_prev, ts)
        xk = np.array(xk_next).flatten()
        
        # 获取系统输出（位置）
        y[k] = xk[0]
        
        # 计算误差（考虑噪声）
        e[k] = yd[k] - (y[k] + noise[k])
        
        # PID控制计算
        u[k] = pid_controller.compute_control(e[k], ts)
        
        # 计算旋翼推力
        para[k] = system.set_rotor_input(u[k])
        
    
    # 绘制结果
    plt.figure()
    
    # 位置跟踪图
    plt.subplot(3, 1, 1)
    plt.plot(time, yd, 'r', linewidth=2, label='Desired')
    plt.plot(time, y, 'k', linewidth=2, label='Estimated')
    plt.xlabel('Time (s)')
    plt.ylabel('Rate (rad/s)')
    plt.title('Rate Tracking')
    plt.legend()
    plt.grid(True)
    
    # 误差图
    plt.subplot(3, 1, 2)
    plt.plot(time, yd - y, 'r', linewidth=2)
    plt.xlabel('Time (s)')
    plt.ylabel('Error (rad/s)')
    plt.title('Tracking Error')
    plt.grid(True)
    
    # 控制输入图
    plt.subplot(3, 1, 3)
    plt.plot(time, u, 'r', linewidth=2)
    plt.xlabel('Time (s)')
    plt.ylabel('Control Input')
    plt.title('Control Input')
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()
    
    # 打印统计信息
    # final_error = abs(yd[-1] - y[-1])
    # max_error = np.max(np.abs(yd - y))
    # print(f"\n仿真统计信息:")
    # print(f"  最终误差: {final_error:.6f}")
    # print(f"  最大误差: {max_error:.6f}")
    # print(f"  平均控制输入: {np.mean(np.abs(u)):.6f}")


if __name__ == "__main__":
    main()